# -*- coding: utf-8 -*-
# @Time    : 2025/9/19 16:04
# @Author  : chenmh
# @File    : dataset.py
# @Desc: 构造数据集

import numpy as np
import torch, typing, datetime, random
from config import bst_config, BSTConfig
from torch import nn, optim
from torch.utils.data import Dataset, DataLoader
import pandas as pd


class BSTDataset(Dataset):

    def __init__(self, config: BSTConfig, num_users: int = 1000):
        self.config = config
        self._num_users = num_users
        self.data = self._generate_virtual_data()

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx) -> dict:
        val = self.data.iloc[idx].to_dict()
        item = val["item"]
        times = val["times"]
        static_vals = []
        for i in range(len(self.config.static_embeddings_vocab_size)):
            static_vals.append(val[f"static_{i}"])
        label = val["label"]
        return {
            "item": torch.tensor(item, dtype=torch.long),
            "times": torch.tensor(times, dtype=torch.float32),
            "static_vals": torch.tensor(static_vals, dtype=torch.long),
            "label": torch.tensor(label, dtype=torch.long),
        }

    def _apply_item(self, x: pd.Series, col_name: str) -> pd.Series:
        """

        :param x: Row[idx]
        :param col_name: 每一个 Row 需要修改的字段
        :return:
        """
        x[col_name] = [0] * x["idx"] + x[col_name][x["idx"]:]
        return x[col_name]

    def _generate_timestamps_ascending(self, count: int = 200) -> list[int]:
        """
        生成近一年内的指定数量的时间戳，并按升序返回。
        参数:
            count (int): 要生成的时间戳数量，默认为 200
        返回:
            list: 升序排列的 Unix 时间戳列表（单位：秒）
        """
        # 获取当前时间
        now = datetime.datetime.now()
        # 计算一年前的时间（近似为365天）
        one_year_ago = now - datetime.timedelta(days=365)
        # 将两个时间转换为 Unix 时间戳（浮点数）
        start_ts = one_year_ago.timestamp()
        end_ts = now.timestamp()
        # 随机生成指定数量的时间戳
        timestamps = [random.uniform(start_ts, end_ts) for _ in range(count)]
        # 按升序排序
        timestamps.sort()
        timestamps = [int(ts) for ts in timestamps]
        return timestamps

    def _generate_virtual_data(self) -> pd.DataFrame:
        """
        生成虚拟的测试数据
        :return:
        """
        df = pd.DataFrame()
        df["item"] = np.random.randint(0, self.config.vocab_size, size=(self._num_users, self.config.seq_len)).tolist()
        df["idx"] = np.random.randint(0, self.config.vocab_size // 2, size=self._num_users).tolist()
        df["item"] = df.apply(lambda x: self._apply_item(x, "item"), axis=1)
        df["times"] = [self._generate_timestamps_ascending(count=self.config.seq_len) for _ in range(self._num_users)]
        df["times"] = df.apply(lambda x: self._apply_item(x, "times"), axis=1)
        for i in range(len(self.config.static_embeddings_vocab_size)):
            df[f"static_{i}"] = np.random.randint(0, self.config.static_embeddings_vocab_size[i],
                                                  size=self._num_users).tolist()
        df["label"] = np.random.choice([0, 1], p=[0.6, 0.4], size=self._num_users)
        del df["idx"]
        return df.reset_index(drop=True)


# bst_dataset = BSTDataset(config=bst_config, num_users=100)
# print(bst_dataset)

train_loader = DataLoader(dataset=BSTDataset(config=bst_config, num_users=1000), batch_size=bst_config.batch_size,
                          shuffle=True)
valid_loader = DataLoader(dataset=BSTDataset(config=bst_config, num_users=300), batch_size=bst_config.batch_size,
                          shuffle=False)
test_loader = DataLoader(dataset=BSTDataset(config=bst_config, num_users=300), batch_size=bst_config.batch_size,
                         shuffle=False)

# if __name__ == '__main__':
#     for batch in train_loader:
#         print(batch)
#         break
